"""
SAM 模型下载脚本
自动下载 Segment Anything 模型文件
"""

import urllib.request
import os
from pathlib import Path
import sys


# 模型配置
MODELS = {
    "vit_h": {
        "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
        "filename": "sam_vit_h_4b8939.pth",
        "size": "2.4GB",
        "description": "最大、最准确（推荐）"
    },
    "vit_l": {
        "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
        "filename": "sam_vit_l_0b3195.pth",
        "size": "1.2GB",
        "description": "中等大小"
    },
    "vit_b": {
        "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
        "filename": "sam_vit_b_01ec64.pth",
        "size": "375MB",
        "description": "最小、最快"
    }
}


def download_file(url, filename):
    """下载文件并显示进度"""
    def show_progress(block_num, block_size, total_size):
        downloaded = block_num * block_size
        percent = min(downloaded * 100.0 / total_size, 100)
        sys.stdout.write(f"\r下载进度: {percent:.1f}% ({downloaded / 1024 / 1024:.1f}MB / {total_size / 1024 / 1024:.1f}MB)")
        sys.stdout.flush()
    
    print(f"开始下载: {filename}")
    print(f"URL: {url}")
    
    try:
        urllib.request.urlretrieve(url, filename, show_progress)
        print(f"\n✓ 下载完成: {filename}")
        return True
    except Exception as e:
        print(f"\n✗ 下载失败: {str(e)}")
        return False


def main():
    """主函数"""
    print("=" * 70)
    print("SAM 模型下载工具")
    print("=" * 70)
    print()
    
    # 创建 models 目录
    models_dir = Path("models")
    models_dir.mkdir(exist_ok=True)
    
    # 显示可用模型
    print("可用的模型:")
    for key, model in MODELS.items():
        print(f"  [{key}] {model['description']} - {model['size']}")
    print()
    
    # 选择模型
    choice = input("请选择要下载的模型 (vit_h/vit_l/vit_b，直接回车下载 vit_h): ").strip().lower()
    
    if not choice:
        choice = "vit_h"
    
    if choice not in MODELS:
        print(f"错误: 无效的选择 '{choice}'")
        return
    
    model = MODELS[choice]
    output_path = models_dir / model["filename"]
    
    # 检查文件是否已存在
    if output_path.exists():
        print(f"文件已存在: {output_path}")
        overwrite = input("是否重新下载? (y/n): ").strip().lower()
        if overwrite != 'y':
            print("取消下载")
            return
        os.remove(output_path)
    
    print()
    print(f"正在下载 {choice} 模型...")
    print(f"文件大小: {model['size']}")
    print(f"保存位置: {output_path}")
    print()
    
    # 下载文件
    success = download_file(model["url"], str(output_path))
    
    if success:
        print()
        print("=" * 70)
        print("✓ 下载成功！")
        print(f"模型文件位置: {output_path.absolute()}")
        print()
        print("接下来您可以:")
        print("1. 运行 python sam_to_psd.py 启动 GUI 工具")
        print("2. 点击【加载模型】按钮，选择刚下载的模型文件")
        print("=" * 70)
    else:
        print()
        print("=" * 70)
        print("✗ 下载失败！")
        print("请检查网络连接，或手动从以下地址下载:")
        print(f"  {model['url']}")
        print(f"然后保存到: {output_path.absolute()}")
        print("=" * 70)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n下载已取消")
    except Exception as e:
        print(f"\n错误: {str(e)}")

